#!/usr/binn/env/python
from random import shuffle, choice, sample, random, randint
from bisect import bisect
from itertools import izip_longest
from math import sqrt

# helper stuff
fst = lambda s: s[0]
snd = lambda s: s[1]


def grouper(n, iterable, fillvalue=None):
    "Collect data into fixed-length chunks or blocks"
    # grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx
    args = [iter(iterable)] * n
    return izip_longest(fillvalue=fillvalue, *args)


def rand_index_pair(start, end):
    """returns two distinct, randomly chosen elements in [start, end]"""
    x0 = randint(start, end)
    x1 = x0
    while x1 == x0:
        x1 = randint(start, end)
    if x0 > x1:
        return (x1, x0)
    return (x0, x1)


def cdf(weights):
    total = sum(weights)
    result = []
    cumsum = 0.0
    for w in weights:
        cumsum += w
        result.append(cumsum / total)
    return result


def prob_choice(population, cdf):
    assert len(population) == len(cdf)
    x = random()
    idx = bisect(cdf, x)
    return population[idx]


##### selectors #####

def tournament_selection(population, num):
    selector = lambda (a, b): a if snd(a) <= snd(b) else b
    for _ in xrange(num):
        yield selector(sample(population, 2))


def rank_selection(population, num):
    population = sorted(population, key=snd)
    N = len(population)
    probs = cdf([(2.0 / N) * (1.0 - (k - 1.0) / (N - 1.0))
                for k in xrange(1, N + 1)])
    for _ in xrange(num):
        yield prob_choice(population, probs)


def roulette_selection(population, num):
    probs = cdf(map(snd, population))
    for _ in xrange(num):
        yield prob_choice(population, probs)


##### cross-overs #####
def one_point_crossover(dad, mom):
    def cmp_snd(a, b):
        return cmp(snd(a), snd(b))

    def cmp_fst(a, b):
        return cmp(fst(a), fst(b))

    cut = randint(0, len(mom))
    # cut both genomes in half
    #print "crossover at {0}".format(cut)
    #print "{0} || {1}".format(" ".join(map(str, dad[:cut])), " ".join(map(str, dad[cut:])))
    #print "{0} || {1}".format(" ".join(map(str, mom[:cut])), " ".join(map(str, mom[cut:])))
    #print "-" * (4 + 2 * (len(mom) - 1))
    dad_set = set(dad[cut:])
    mom_set = set(mom[cut:])
    child0 = dad[:cut] + [x for x in mom if x in dad_set]
    child1 = mom[:cut] + [x for x in dad if x in mom_set]
    #print "{0} || {1}".format(" ".join(map(str, dad[:cut])), " ".join(map(str, dad[cut:])))
    #print "{0} || {1}".format(" ".join(map(str, mom[:cut])), " ".join(map(str, mom[cut:])))
    #print
    if random() < 0.5:
        return child0
    else:
        return child1


def partially_matched_crossover(dad, mom):
    a, b = rand_index_pair(0, len(dad))
    # cut both genomes
    #dad_0, dad_m, dad_1 = dad[:a], dad[a:b], dad[b:]
    #mom_0, mom_m, mom_1 = mom[:a], mom[a:b], mom[b:]
    #print "crossover at {0}, {1}".format(a, b)
    #print "{0} || {1} || {2}".format(" ".join(map(str, dad_0)), " ".join(map(str, dad_m)), " ".join(map(str, dad_1)))
    #print "{0} || {1} || {2}".format(" ".join(map(str, mom_0)), " ".join(map(str, mom_m)), " ".join(map(str, mom_1)))
    #print "-" * (6 + 2 * (len(mom) - 1))

    child0 = dad[:]
    child1 = mom[:]

    p1, p2 = [0] * len(dad), [0] * len(dad)
    for i in xrange(len(dad)):
        p1[child0[i]] = i
        p2[child1[i]] = i

    for i in xrange(a, b):
        # Keep track of the selected values
        temp1 = child0[i]
        temp2 = child1[i]
        # Swap the matched value
        child0[i], child0[p1[temp2]] = temp2, temp1
        child1[i], child1[p2[temp1]] = temp1, temp2
        # Position bookkeeping
        p1[temp1], p1[temp2] = p1[temp2], p1[temp1]
        p2[temp1], p2[temp2] = p2[temp2], p2[temp1]

    #dad = dad[:cut] + [x for x in mom if x in dad_set]
    #mom = mom[:cut] + [x for x in dad if x in mom_set]
    #print "{0} || {1} || {2}".format(" ".join(map(str, child0[:a])), " ".join(map(str, child0[a:b])), " ".join(map(str, child0[b:])))
    #print "{0} || {1} || {2}".format(" ".join(map(str, child1[:a])), " ".join(map(str, child1[a:b])), " ".join(map(str, child1[b:])))
    assert(len(set(child0)) == len(set(child1)) == len(mom) == len(dad))
    if random() < 0.5:
        return child0
    else:
        return child1


def edge_recombination_crossover(ind0, ind1):
    neighbors = [set() for _ in xrange(len(ind0))]
    for u, v in zip(ind0, ind0[1:]) + [(ind0[-1], ind0[0])]:
        neighbors[u].add(v)
        neighbors[v].add(u)
    for u, v in zip(ind1, ind1[1:]) + [(ind1[-1], ind1[0])]:
        neighbors[u].add(v)
        neighbors[v].add(u)

    if random() < 0.5:
        current = ind0[0]
    else:
        current = ind1[0]

    child = [current]
    seen = set([current])
    while len(child) < len(ind0):
        for n in neighbors:
            if current in n:
                n.remove(current)
        smallest = None
        smallest_n = None
        for n in neighbors[current]:
            if len(neighbors[n]) < smallest_n or smallest_n == None:
                smallest = n
                smallest_n = len(neighbors[n])
        current = smallest
        if current == None:
            current = choice([i for i in xrange(len(ind0)) if i not in seen])
        child.append(current)
        seen.add(current)
    assert(len(child) == len(ind0))
    return child


##### mutators #####
def swap_mutator(xs):
    x0, x1 = rand_index_pair(0, len(xs) - 1)
    xs[x0], xs[x1] = xs[x1], xs[x0]
    return xs


def inversion_mutator(xs):
    x0, x1 = rand_index_pair(0, len(xs) - 1)
    xs[x0:x1 + 1] = list(reversed(xs[x0:x1 + 1]))
    return xs


##### analyze #####
def analyze_population(iteration, maxiter, population, children):
    min_ind = min(population, key=snd)[1]
    max_ind = max(population, key=snd)[1]
    mean = sum(map(snd, population)) / float(len(population))
    variance = sum((x[1] - mean)**2 for x in population) / float(len(population) - 1)
    print "{0:{2}d}/{1} ({3}), min={4}, max={5}, mean={6:.2f}, dev={7:.2f} ({8} unique childs, {9} childs + parents)".format(
        iteration, maxiter, len(str(maxiter)),
        len(population), min_ind, max_ind, mean, sqrt(variance), len(set(tuple(xs) for xs in children)), 
        len(set([tuple(xs) for xs in children])) + len(set([tuple(xs) for (xs, _) in population])))


def tspEA(instance,
          crossover_selection,
          crossover,
          mutation,
          ngen_selection,
          num_parents,
          num_children,
          mutation_rate,
          maxiter,
          elitism):

    # initialize population
    population = []
    for _ in xrange(num_parents):
        xs = range(len(instance))
        shuffle(xs)
        population.append((xs, instance.tour_length(xs)))

    # start loop
    best_individuum = min(population, key=snd)
    print (0, best_individuum)
    for iteration in xrange(1, maxiter + 1):
        old_best = best_individuum
        #print "iteration {0}".format(iteration)
        tmp_children = []
        # select pairs for crossover + children
        for ind0, ind1 in grouper(2, crossover_selection(population, 2 * num_children)):
            child = crossover(ind0[0], ind1[0])
            tmp_children.append(child)

        # mutate
        for child in tmp_children:
            if random() < mutation_rate:
                mutation(child)

        #analyze_population(iteration, maxiter, population, tmp_children)

        # evaluation
        for child in tmp_children:
            child_fitness = instance.tour_length(child)
            individuum = (child, child_fitness)
            if child_fitness < best_individuum[1]:
                best_individuum = individuum
            population.append((child, child_fitness))

        #print len(population)

        # select new population
        new_population = []
        seen_individuums = set()
        if elitism is True:
            new_population = [best_individuum]
            seen_individuums.add(tuple(best_individuum[0]))
        tries = 0
        while len(new_population) < num_parents:
            if tries >= 20:
                break
            tries += 1
            for individuum in ngen_selection(population, 2 * num_parents):
                if len(new_population) >= num_parents:
                    break
                hashable = tuple(individuum[0])
                if hashable not in seen_individuums:
                    seen_individuums.add(hashable)
                    new_population.append(individuum)
        population = new_population
        if best_individuum != old_best:
            print (iteration, best_individuum)
    print best_individuum


if __name__ == '__main__':
    import tsp
    import argparse

    instance = tsp.TSP.from_file("data/TSP/ulysses16edit.tsp")
    tspEA(instance=instance,
          crossover_selection=tournament_selection,
          crossover=one_point_crossover,
          mutation=swap_mutator,
          ngen_selection=tournament_selection,
          num_parents=50,
          num_children=50,
          mutation_rate=0.05,
          maxiter=1000,
          elitism=True)
